package ocr

import (
	"log"
	"path/filepath"
	"time"

	"github.com/sirupsen/logrus"
	"gocv.io/x/gocv"
)

type TextRecognizer struct {
	*PaddleModel
	batchNum int
	textLen  int
	shape    []int
	charType string
	labels   []string
}

func NewTextRecognizer(modelSource string, config *Config) *TextRecognizer {
	shapes := config.RecImageShape //[]int{3, 32, 320}
	// if v, ok := args["rec_image_shape"]; ok {
	// 	for i, s := range v.([]interface{}) {
	// 		shapes[i] = s.(int)
	// 	}
	// }
	labelpath := config.RecCharDictPath
	labels := readLines2StringSlice(labelpath)
	if config.UseSpaceChar {
		labels = append(labels, " ")
	}
	rec := &TextRecognizer{
		PaddleModel: NewPaddleModel("rec", config),
		batchNum:    config.RecBatchNum,   //getInt(args, "rec_batch_num", 10),
		textLen:     config.MaxTextLength, //getInt(args, "max_text_length", 25),
		charType:    config.RecCharType,   //getString(args, "rec_char_type", "ch"),
		shape:       shapes,
		labels:      labels,
	}
	if checkModelExists(modelSource) {
		modelSource, _ = downloadModel(filepath.Join(config.ModleDir, "rec/ch"), modelSource)
	} else {
		log.Panicf("rec model path: %v not exist! Please check!", modelSource)
	}
	rec.LoadModel(modelSource)
	return rec
}

func (rec *TextRecognizer) Run(imgs []gocv.Mat, bboxes [][][]int) []OCRText {
	recResult := make([]OCRText, 0, len(imgs))
	batch := rec.batchNum
	var recTime int64 = 0
	c, h, w := rec.shape[0], rec.shape[1], rec.shape[2]
	for i := 0; i < len(imgs); i += batch {
		logrus.WithField("i", i).WithField("len(imgs)", len(imgs)).WithField("batch", batch).Info("recognition loop")
		j := i + batch
		if len(imgs) < j {
			j = len(imgs)
		}

		maxwhratio := 0.0
		for k := i; k < j; k++ {
			h, w := imgs[k].Rows(), imgs[k].Cols()
			ratio := float64(w) / float64(h)
			if ratio > maxwhratio {
				maxwhratio = ratio
			}
		}

		if rec.charType == "ch" {
			w = int(32 * maxwhratio)
		}
		normimgs := make([]float32, (j-i)*c*h*w)

		for k := i; k < j; k++ {
			data := crnnPreprocess(imgs[k], rec.shape, []float32{0.5, 0.5, 0.5},
				[]float32{0.5, 0.5, 0.5}, 255.0, maxwhratio, rec.charType)
			copy(normimgs[(k-i)*c*h*w:], data)
		}

		st := time.Now()

		rec.inputs[rec.inNames[0]].Reshape([]int32{int32(j - i), int32(c), int32(h), int32(w)})
		rec.inputs[rec.inNames[0]].CopyFromCpu(normimgs)

		rec.predictor.Run()

		// rec.outputs[rec.outNames[0]].Reshape([]int32{int32(j - i), int32(rec.textLen), int32(len(rec.labels))})

		var outShape = rec.outputs[rec.outNames[0]].Shape()

		predictBatch := make([]float32, numElements(rec.outputs[rec.outNames[0]].Shape()))
		rec.outputs[rec.outNames[0]].CopyToCpu(predictBatch)
		predictLod := rec.outputs[rec.outNames[0]].Lod()

		logrus.WithField("predict batch", len(predictBatch)).WithField("len of rec idx", len(predictLod)).Info("recognition loop")
		// recIdxBatch := make([]int64, numElements(rec.outputs[rec.outNames[1]].Shape()))
		// rec.outputs[rec.outNames[1]].CopyToCpu(recIdxBatch)
		// recIdxLod := rec.outputs[rec.outNames[1]].Lod()

		recTime += int64(time.Since(st).Milliseconds())

		for m := int32(0); m < outShape[0]; m++ {
			words := ""
			score := 0.0
			count := 0
			for n := int32(0); n < outShape[1]; n++ {
				var offset = m*outShape[1]*outShape[2] + n*outShape[2]
				var size = outShape[2]
				argMaxID, maxVal := argmax(predictBatch[offset : offset+size])
				if argMaxID > 0 {
					// fmt.Println("argMaxID: ", argMaxID, "maxVal: ", maxVal, "label: ", rec.labels[argMaxID])
					words += rec.labels[argMaxID-1]
					score += float64(maxVal)
					count++
				}
			}
			recResult = append(recResult, OCRText{
				BBox:  bboxes[i+int(m)],
				Text:  words,
				Score: score,
			})
		}

	}
	logrus.WithField("rec_time", recTime).Info("recognition time")
	return recResult
}
